import binascii
from gmssl.sm4 import CryptSM4, SM4_DECRYPT, SM4_ENCRYPT

"""SM4工具"""
# 初始化
crypt_sm4 = CryptSM4()

def encrypt_ecb(msg, key):
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    encrypted = crypt_sm4.crypt_ecb(msg) #  bytes类型
    print('encrypt_value: ', encrypted)
    print('enc_hex: ', encrypted.hex())
    return encrypted

def decrypt_ecb(enc_msg, key):
    crypt_sm4.set_key(key, SM4_DECRYPT)
    decrypted = crypt_sm4.crypt_ecb(enc_msg) #  bytes类型
    print('decrypted: ', decrypted)
    return decrypted

def encrypt_ecb_to_hex(msg: str, hex_key: str):
    key = bytes.fromhex(hex_key)
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    encrypted = crypt_sm4.crypt_ecb(msg.encode('utf-8')) #  bytes类型
    #print('encrypt_value: ', encrypted)
    #print('enc_hex: ', encrypted.hex())
    return encrypted.hex()

def decrypt_ecb_from_hex(encrypted_hex: str, hex_key: str):
    key = bytes.fromhex(hex_key)
    enc_bytes = binascii.unhexlify(encrypted_hex)
    crypt_sm4.set_key(key, SM4_DECRYPT)
    decrypted = crypt_sm4.crypt_ecb(enc_bytes) #  bytes类型
    #print('decrypted: ', decrypted)
    return decrypted.decode('utf-8')

def encrypt_cbc(msg, key, iv=None):
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    encrypted = crypt_sm4.crypt_cbc(iv , msg) #  bytes类型
    print('encrypt_value: ', encrypted)
    print('enc_hex: ', encrypted.hex())
    return encrypted

def decrypt_ecb(enc_msg, key):
    crypt_sm4.set_key(key, SM4_DECRYPT)
    decrypted = crypt_sm4.crypt_ecb(enc_msg) #  bytes类型
    #print('decrypted: ', decrypted)
    return decrypted

def decrypt_cbc(enc_msg, key, iv=None):
    crypt_sm4.set_key(key, SM4_DECRYPT)
    decrypted = crypt_sm4.crypt_cbc(iv , enc_msg) #  bytes类型
    print('decrypted: ', decrypted)
    return decrypted

if __name__ == '__main__':
    #hex_key = '0123456789abcdeffedcba9876543210'
    hex_key = '8a81609a885ae99a70f399548f34c2f5'
    key = bytes.fromhex(hex_key)
    msg = '123456afsfsfs'.encode()
    #msg = b'123456afsfsfs'
    #encrypted = encrypt_ecb(msg, key)
    #print('encrypted: ', encrypted)
    #encrypted_hex = encrypt_ecb_to_hex(msg, key)
    #print('encrypted_hex: ', encrypted_hex)
    encrypted_hex = '283CD0C569FA60C75EC5D6E2674C081F41D45DD906AFA2BE79556289D6E73A4A'.lower()
    byte_str = binascii.unhexlify(encrypted_hex)
    print('encrypted_bytes: ', byte_str)
    
    decrypted_bytes = decrypt_ecb(byte_str, key)
    print('decrypted_bytes: ', decrypted_bytes)
    decrypted_str = decrypted_bytes.decode("utf-8")
    print('decrypted_str: ', decrypted_str)
